"""
统一客户端类
"""

from rich.console import Console
from rich.prompt import Prompt
from rich.panel import Panel
from rich.table import Table
from agent_mode import AgentMode
from control_mode import ControlMode
from relay_api import RelayAPI

console = Console()

class SSHRemoteClient:
    def __init__(self):
        self.agent_mode = AgentMode()
        self.control_mode = ControlMode()
        self.api = RelayAPI()
    
    def run(self):
        """运行客户端主程序"""
        self._show_banner()
        
        while True:
            choice = self._show_main_menu()
            
            if choice == "1":
                self._run_agent_mode()
            elif choice == "2":
                self._run_control_mode()
            elif choice == "3":
                console.print("[bold green]感谢使用SSH远程控制工具！[/bold green]")
                break
            else:
                console.print("[bold red]无效选择，请重新输入[/bold red]")
    
    def _show_banner(self):
        """显示程序横幅"""
        banner = """
[bold cyan]
╔═══════════════════════════════════════╗
║        SSH 远程控制工具 v1.0          ║
║     支持内网穿透的安全远程连接        ║
╚═══════════════════════════════════════╝
[/bold cyan]
        """.strip()
        
        console.print(banner)
        console.print()
    
    def _show_main_menu(self) -> str:
        """显示主菜单"""
        menu_text = """
[bold white]请选择工作模式:[/bold white]

[bold green]1.[/bold green] 被控模式 - 等待他人连接
[bold blue]2.[/bold blue] 控制模式 - 连接他人
[bold red]3.[/bold red] 退出程序
        """.strip()
        
        panel = Panel(
            menu_text,
            title="[bold cyan]主菜单[/bold cyan]",
            border_style="cyan",
            padding=(1, 2)
        )
        
        console.print(panel)
        
        choice = Prompt.ask(
            "[bold cyan]请输入选项 (1-3)[/bold cyan]",
            choices=["1", "2", "3"],
            default="1"
        )
        
        return choice
    
    def _run_agent_mode(self):
        """运行被控模式"""
        try:
            self.agent_mode.start()
        except Exception as e:
            console.print(f"[bold red]被控模式运行错误: {str(e)}[/bold red]")
        
        console.print("\n按 [bold yellow]Enter[/bold yellow] 返回主菜单...")
        input()
    
    def _run_control_mode(self):
        """运行控制模式"""
        try:
            self.control_mode.start()
        except Exception as e:
            console.print(f"[bold red]控制模式运行错误: {str(e)}[/bold red]")
        
        console.print("\n按 [bold yellow]Enter[/bold yellow] 返回主菜单...")
        input()





if __name__ == '__main__':
    SSHRemoteClient().run()


